from . import face_model
# import argparse
import numpy as np
from scipy import spatial
import pandas as pd
from sklearn import neighbors
import datetime
unknow_name = 'Unknown'
similar_cost = 0.8

import time

# parser = argparse.ArgumentParser(description='face model test')
# # general
# parser.add_argument('--image-size', default='112,112', help='')
# parser.add_argument('--model', default='face_recog/models/model-r34-amf/model,0', help='path to load model.')
# parser.add_argument('--ga-model', default='/Users/sshuair/AI/face-recognition/insightface/models/gamodel-r50/model,0', help='path to load model.')
# parser.add_argument('--gpu', default=-1, type=int, help='gpu id')
# parser.add_argument('--det', default=0, type=int, help='mtcnn option, 1 means using R+O, 0 means detect from begining')
# parser.add_argument('--flip', default=0, type=int, help='whether do lr flip aug')
# parser.add_argument('--threshold', default=1.24, type=float, help='ver dist threshold')
# args = parser.parse_args()

args={}
args['image_size'] = '112,112'
# args['model'] = 'face_recog/models/model-r34-amf/model,0'
args['model'] = 'face_recog/models/model-mobile/model,0'
args['gpu']=-1
args['det']=0
args['flip']=0
args['threshold'] = 1.24

# 加载人脸检测和特征提取模型
model = face_model.FaceModel(args)
# 加载人脸特征数据库，构建K-D Tree
face_df = pd.read_csv('face_recog/deploy/face_feat_mobile_v2.txt', header=None)
face_df.columns = face_df.columns.map(str)
X = face_df.drop(labels=['0','1'], axis=1)
names = face_df['0']
user_ids = face_df['1']
kdtree = neighbors.KDTree(X)


def search_name(feat):
    """匹配已有人脸数据库，两个条件判断是否是一个人，首先筛选出cos距离大于0.33的值，然后计算k近邻
    1. 查询最近的一个
    2. 计算cos值，如果大于0.33则认为是同一个人，否是是unknown

    Args:
        feat ([type]): [description]
    """
    distances, indices = kdtree.query(np.array(feat).reshape(-1, len(feat)))
    indices = indices[0][0]
    nearst_feat  = X.loc[indices,:].tolist()
    cos_dist = spatial.distance.cosine(nearst_feat, feat)
    if cos_dist < similar_cost:
        return names[indices], user_ids[indices]
    else:
        return 'Unknown', None

def face_recog(frame):
    """
    Args:
        img ():
    """
    result = []
    st=time.time()
    img_faces = model.get_input2(frame) #返回 boxes, points, face_aligned
    # print(img_faces)
    print('detection: ', time.time()-st)

    st=time.time()
    if img_faces is None:
        return None

    for bbox, keypoint, face_aligned in zip(*img_faces):
        print(bbox.shape)
        # if int(bbox[0]) > 0  and int(bbox[2]) > 0 and int(bbox[1])>0 and int(bbox[3])> 0:
        #     print(123)
        head_img = frame[int(bbox[1]): int(bbox[3]), int(bbox[0]): int(bbox[2])]
        print(head_img.shape)
        if len(head_img) > 0:
            image_feature = model.get_feature(face_aligned)
            print('image_feature', image_feature.shape)
            name, user_id = search_name(image_feature)
            if user_id is not None:
                user_id = int(user_id)
            result.append({'name':name,'user_id': user_id,'image_feature':image_feature.tolist(), 'bbox':list(bbox[0:4])})

    return result

if __name__ == "__main__":
    from PIL import Image
    img = np.array(Image.open('face_recog/deploy/IMG_6382.jpg'))
    print(face_recog(img))
